import org.apache.spark.mllib.regression.LabeledPoint
import org.apache.spark.rdd.RDD
import org.apache.spark.{SparkConf, SparkContext}

import org.apache.spark.mllib.tree.DecisionTree
import org.apache.spark.mllib.tree.model.DecisionTreeModel
import org.apache.spark.mllib.util.MLUtils


object DecisionTreeDemo {

  def main(args: Array[String]): Unit = {
    val sc = new SparkContext(new SparkConf().setMaster("local[*]").setAppName("DecisionTreeDemo"))

    val data: RDD[LabeledPoint] = MLUtils.loadLibSVMFile(sc, "作业/resources/example/sample_libsvm_data.txt")
    val splits: Array[RDD[LabeledPoint]] = data.randomSplit(Array(0.7, 0.3))
    val (trainingData, testData) = (splits(0), splits(1))
    val model: DecisionTreeModel = DecisionTree.trainClassifier(trainingData, numClasses = 2, categoricalFeaturesInfo = Map[Int, Int](),
      impurity = "gini", maxDepth = 5, maxBins = 64)


    val label_Preds: RDD[(Double, Double)] = testData.map { point =>
      val prediction: Double = model.predict(point.features)
      (point.label, prediction)
    }
    val testErr: Double = label_Preds.filter(r => r._1 != r._2).count().toDouble / testData.count()
    println(s"Test Error = $testErr")
    println(s"Learned classification tree model:\n ${model.toDebugString}")

//    model.save(sc, "./output/myDecisionTreeClassificationModel")
//    val sameModel: DecisionTreeModel = DecisionTreeModel.load(sc, "./output/myDecisionTreeClassificationModel")

    sc.stop()
  }
}

// scalastyle:on println
